
import sys
import os

# 假设你的包位于"/path/to/your/package"目录
package_path = "/home/xjz/catkin_ws/src/my_get_model_state_pkg/scripts/"
if package_path not in sys.path:
    sys.path.append(package_path)

from  get_aliengo_state import model_state
from  get_dof_state import dof_state
from  onnx_model import onnx_model
from  publish_aliengo_state import publisher
import rospy
import numpy as np
import time

def get_axis_params(value, axis_idx, x_value=0., dtype=np.float64, n_dims=3):
    """construct arguments to `Vec` according to axis index.
    """
    zs = np.zeros((n_dims,))
    assert axis_idx < n_dims, "the axis dim should be within the vector dimensions"
    zs[axis_idx] = 1.
    params = np.where(zs == 1., value, zs)
    params[0] = x_value
    return np.array(list(params.astype(dtype)))

import numpy as np

def quat_rotate(q, v):
    shape = q.shape
    q_w = q[:, -1]
    q_vec = q[:, :3]
    
    # 对应a的部分
    a = v * (2.0 * q_w ** 2 - 1.0)[:, np.newaxis]
    
    # 对应b的部分，使用np.cross进行叉乘
    b = np.cross(q_vec, v, axis=-1) * (q_w[:, np.newaxis] * 2.0)
    
    # 对应c的部分，这里利用np.outer和np.sum来模拟torch.bmm的行为
    dot_product = np.sum(q_vec * v, axis=-1)[:, np.newaxis]
    c = q_vec * dot_product * 2.0
    
    return a + b + c

def quat_rotate_inverse(q, v):
    shape = q.shape
    q_w = q[:, -1]
    q_vec = q[:, :3]
    
    # 对应a的部分
    a = v * (2.0 * q_w ** 2 - 1.0)[:, np.newaxis]
    
    # 对应b的部分，使用np.cross进行叉乘
    b = np.cross(q_vec, v, axis=-1) * (q_w[:, np.newaxis] * 2.0)
    
    # 对应c的部分，这里使用np.outer代替torch.bmm
    c = q_vec * np.sum(q_vec * v, axis=-1)[:, np.newaxis] * 2.0
    
    return a - b + c


class  aliengo:
    def __init__(self) -> None:
        self.init_config()
        rospy.init_node('aliengo_run', anonymous=True)
        self.model_state_geter=model_state(model_name = 'aliengo_gazebo',command="/gazebo/get_model_state")
        self.dof_state_geter=dof_state('/gazebo/get_joint_properties',self.joint_name_list)
        self.joint_publisher = publisher(self.command_topics) 
        self.model=onnx_model(self.onnx_model_path)
    
    def init_config(self):
            self.joint_name_order=[
                        'FL_hip_joint', 
                        'FL_thigh_joint',
                        'FL_calf_joint', 
                        'FR_hip_joint', 
                        'FR_thigh_joint', 
                        'FR_calf_joint', 
                        'RL_hip_joint', 
                        'RL_thigh_joint', 
                        'RL_calf_joint', 
                        'RR_hip_joint', 
                        'RR_thigh_joint', 
                        'RR_calf_joint']
            self.command_topics = [
                "/aliengo_gazebo/FL_hip_joint/command",
                "/aliengo_gazebo/FL_thigh_joint/command",                
                "/aliengo_gazebo/FL_calf_joint/command",
                "/aliengo_gazebo/FR_hip_joint/command",
                "/aliengo_gazebo/FR_thigh_joint/command",
                "/aliengo_gazebo/FR_calf_joint/command",
                "/aliengo_gazebo/RL_hip_joint/command",
                "/aliengo_gazebo/RL_thigh_joint/command",
                "/aliengo_gazebo/RL_calf_joint/command",
                "/aliengo_gazebo/RR_hip_joint/command",
                "/aliengo_gazebo/RR_thigh_joint/command",
                "/aliengo_gazebo/RR_calf_joint/command"
                ]
            defaultJointAngles={  # = target angles when action = 0.0
                        'FL_hip_joint': 0.1,   # [rad]
                        'RL_hip_joint': 0.1 ,  # [rad]
                        'FR_hip_joint': -0.1 , # [rad]
                        'RR_hip_joint': -0.1  , # [rad]
                        'FL_thigh_joint': 0.8  ,   # [rad]
                        'RL_thigh_joint': 1.  , # [rad]
                        'FR_thigh_joint': 0.8 ,    # [rad]
                        'RR_thigh_joint': 1.  , # [rad]
                        'FL_calf_joint': -1.5  ,  # [rad]
                        'RL_calf_joint': -1.5  ,  # [rad]
                        'FR_calf_joint': -1.5 , # [rad]
                        'RR_calf_joint': -1.5  ,  # [rad]
              }
            self.commands=np.array([1.5,0,0],dtype=np.float32)[None,:]
            self.default_dof_pos=[]
            for joint_name in self.joint_name_order:
                 self.default_dof_pos.append(defaultJointAngles[joint_name])
            self.default_dof_pos=np.array(self.default_dof_pos)[None,:]
            self.gravity_vec = np.repeat(get_axis_params(-1., 2)[np.newaxis, :],1,axis=0)
            self.actions = np.zeros((1,12),dtype=np.float32)
            self.lin_vel_scale=2.0
            self.ang_vel_scale=0.25
            self.dof_pos_scale=1.0
            self.dof_vel_scale=0.05
            self.onnx_model_path = '/home/xjz/catkin_ws/src/my_get_model_state_pkg/scripts/models/xjz_aliengo.onnx'  
            self.action_scale=0.5
            self.clip_obs=5.0
            print()
            

    def run(self):
        root_state=self.model_state_geter.get_state()[None, :]
        dof_state=self.dof_state_geter.get_joint_properties()
        dof_pos=dof_state[:,0][None,:]
        dof_vel=dof_state[:,1][None,:]
        obs=self.compute_aliengo_observations(root_state,self.commands, dof_pos,self.default_dof_pos,dof_vel,self.gravity_vec,self.actions,       
                                self.lin_vel_scale,
                                self.ang_vel_scale,
                                self.dof_pos_scale,
                                self.dof_vel_scale)
        obs=obs.astype(np.float32)
        self.actions=self.model.infference(obs)
        # targets = self.action_scale * self.actions + self.default_dof_pos
        targets = self.action_scale * self.actions 
        # targets=self.actions
        self.joint_publisher.publish_topics(targets.squeeze().tolist())
        pass

    def compute_aliengo_observations(self,root_states,
                                commands,
                                dof_pos,
                                default_dof_pos,
                                dof_vel,
                                gravity_vec,
                                actions,
                                lin_vel_scale,
                                ang_vel_scale,
                                dof_pos_scale,
                                dof_vel_scale
                                ):

        # type: (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, float, float, float, float) -> Tensor
        base_quat = root_states[:, 3:7]
        base_lin_vel = quat_rotate_inverse(base_quat, root_states[:, 7:10]) * lin_vel_scale
        base_ang_vel = quat_rotate_inverse(base_quat, root_states[:, 10:13]) * ang_vel_scale
        projected_gravity = quat_rotate(base_quat, gravity_vec)
        dof_pos_scaled = (dof_pos - default_dof_pos) * dof_pos_scale

        commands_scaled = commands*np.array([lin_vel_scale, lin_vel_scale, ang_vel_scale])

        obs =  np.concatenate((base_lin_vel, base_ang_vel,  projected_gravity,  commands_scaled, dof_pos_scaled, dof_vel*dof_vel_scale, actions
                        ), axis=-1)

        return obs




if __name__=="__main__":
    robot=aliengo()
    control_frequency_hz = 60
    # 计算每次循环应该等待的时间（秒）
    loop_sleep_time = 1.0 / control_frequency_hz
    for i in range(1000):
        robot.run()
        time.sleep(loop_sleep_time)
        